import cv2
import numpy as np
from typing import List, Dict, Tuple
import logging

class NodeDetector:
    """流程图节点检测模块"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.node_types = {
            'rectangle': 'process',
            'diamond': 'decision', 
            'circle': 'start_end',
            'ellipse': 'flow'
        }
    
    def detect_nodes(self, image: np.ndarray, binary_image: np.ndarray) -> List[Dict]:
        """检测流程图中的所有节点"""
        self.logger.info("开始检测流程图节点")
        
        nodes = []
        
        # 查找轮廓
        contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        for i, contour in enumerate(contours):
            # 过滤太小的轮廓
            area = cv2.contourArea(contour)
            if area < 100:  # 最小面积阈值
                continue
            
            # 获取边界框
            x, y, w, h = cv2.boundingRect(contour)
            
            # 计算轮廓的近似形状
            epsilon = 0.02 * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            
            # 判断节点类型
            node_type = self._classify_node_shape(approx, w, h)
            
            if node_type:
                node = {
                    'id': f'node_{i}',
                    'type': node_type,
                    'position': {
                        'x': x,
                        'y': y,
                        'width': w,
                        'height': h
                    },
                    'contour': contour,
                    'confidence': self._calculate_confidence(approx, node_type, area)
                }
                nodes.append(node)
        
        # 按置信度排序
        nodes.sort(key=lambda x: x['confidence'], reverse=True)
        
        self.logger.info(f"检测到 {len(nodes)} 个节点")
        return nodes
    
    def _classify_node_shape(self, approx: np.ndarray, width: int, height: int) -> str:
        """根据轮廓形状分类节点类型"""
        num_vertices = len(approx)
        aspect_ratio = width / height if height > 0 else 1
        
        # 圆形节点（开始/结束）
        if num_vertices >= 8 and 0.8 <= aspect_ratio <= 1.2:
            # 计算轮廓的圆形度
            area = cv2.contourArea(approx)
            perimeter = cv2.arcLength(approx, True)
            circularity = 4 * np.pi * area / (perimeter * perimeter) if perimeter > 0 else 0
            
            if circularity > 0.7:
                return 'start_end'  # 后续会根据位置和文本判断具体类型
        
        # 菱形节点（决策）
        if num_vertices == 4:
            # 检查是否为菱形（对角线长度相近）
            points = approx.reshape(-1, 2)
            if len(points) == 4:
                # 计算对角线长度
                diag1 = np.linalg.norm(points[0] - points[2])
                diag2 = np.linalg.norm(points[1] - points[3])
                
                # 对角线长度比例
                diag_ratio = min(diag1, diag2) / max(diag1, diag2) if max(diag1, diag2) > 0 else 0
                
                if diag_ratio > 0.6:  # 对角线长度相近
                    return 'decision'
        
        # 椭圆形节点
        if num_vertices >= 6 and (aspect_ratio > 1.5 or aspect_ratio < 0.7):
            return 'flow'
        
        # 矩形节点（默认）
        if num_vertices == 4:
            return 'process'
        
        return None
    
    def _calculate_confidence(self, approx: np.ndarray, node_type: str, area: float) -> float:
        """计算节点检测的置信度"""
        confidence = 0.5  # 基础置信度
        
        # 根据面积调整置信度
        if 500 <= area <= 5000:
            confidence += 0.2
        elif area > 5000:
            confidence += 0.1
        
        # 根据节点类型调整置信度
        if node_type == 'decision':
            # 菱形检测相对困难，降低置信度
            confidence -= 0.1
        elif node_type == 'start_end':
            # 圆形检测相对准确
            confidence += 0.1
        
        # 根据轮廓复杂度调整置信度
        perimeter = cv2.arcLength(approx, True)
        if perimeter > 0:
            compactness = area / (perimeter * perimeter)
            if 0.05 <= compactness <= 0.2:
                confidence += 0.1
        
        return min(1.0, max(0.0, confidence))
    
    def filter_overlapping_nodes(self, nodes: List[Dict], overlap_threshold: float = 0.3) -> List[Dict]:
        """过滤重叠的节点"""
        filtered_nodes = []
        
        for i, node1 in enumerate(nodes):
            is_overlapping = False
            
            for j, node2 in enumerate(filtered_nodes):
                overlap_ratio = self._calculate_overlap(node1, node2)
                if overlap_ratio > overlap_threshold:
                    # 保留置信度更高的节点
                    if node1['confidence'] > node2['confidence']:
                        filtered_nodes.remove(node2)
                        filtered_nodes.append(node1)
                    is_overlapping = True
                    break
            
            if not is_overlapping:
                filtered_nodes.append(node1)
        
        return filtered_nodes
    
    def _calculate_overlap(self, node1: Dict, node2: Dict) -> float:
        """计算两个节点的重叠比例"""
        pos1 = node1['position']
        pos2 = node2['position']
        
        # 计算交集区域
        x1 = max(pos1['x'], pos2['x'])
        y1 = max(pos1['y'], pos2['y'])
        x2 = min(pos1['x'] + pos1['width'], pos2['x'] + pos2['width'])
        y2 = min(pos1['y'] + pos1['height'], pos2['y'] + pos2['height'])
        
        if x2 <= x1 or y2 <= y1:
            return 0.0
        
        intersection = (x2 - x1) * (y2 - y1)
        area1 = pos1['width'] * pos1['height']
        area2 = pos2['width'] * pos2['height']
        
        # 返回较小的重叠比例
        return min(intersection / area1, intersection / area2)
    
    def detect_nodes_with_filtering(self, image: np.ndarray, binary_image: np.ndarray) -> List[Dict]:
        """检测节点并过滤重叠"""
        nodes = self.detect_nodes(image, binary_image)
        filtered_nodes = self.filter_overlapping_nodes(nodes)
        
        # 后处理节点类型
        filtered_nodes = self._post_process_node_types(filtered_nodes, image)
        
        self.logger.info(f"过滤后剩余 {len(filtered_nodes)} 个节点")
        return filtered_nodes
    
    def _post_process_node_types(self, nodes: List[Dict], image: np.ndarray) -> List[Dict]:
        """后处理节点类型，特别是start_end类型的细化"""
        for node in nodes:
            if node['type'] == 'start_end':
                # 根据位置和文本内容判断是start还是end
                pos = node['position']
                y_position = pos['y']
                image_height = image.shape[0]
                
                # 如果节点在图像上半部分，更可能是start
                if y_position < image_height * 0.3:
                    node['type'] = 'start'
                # 如果节点在图像下半部分，更可能是end
                elif y_position > image_height * 0.7:
                    node['type'] = 'end'
                else:
                    # 根据文本内容判断
                    text = node.get('text', '').lower()
                    if any(keyword in text for keyword in ['开始', 'start', '输入']):
                        node['type'] = 'start'
                    elif any(keyword in text for keyword in ['结束', 'end', '输出', '退出']):
                        node['type'] = 'end'
                    else:
                        # 默认根据位置判断
                        node['type'] = 'start' if y_position < image_height * 0.5 else 'end'
        
        return nodes 